import os
import json


def get_classes2id(dir_json):
    """
    获得所有json对应的label
    :param dir_json: json文件夹
    :return:
        classes2id :json的label对应的txt的label
    """
    classes2id = {}
    jsons = os.listdir(dir_json)
    for i in jsons:
        if os.path.splitext(i)[-1] == ".json":
            json_path = os.path.join(dir_json, i)
            with open(json_path, 'r', encoding="utf-8") as f:
                json_data = json.load(f)
                # print(json_data['shapes'])
                for j in json_data['shapes']:
                    if j['label'] not in classes2id:
                        classes2id[j['label']] = int(j['label'])
    print(classes2id)
    return classes2id


def json2txt(classes2id, path_json, path_txt):  # 可修改生成格式
    with open(path_json, 'r', encoding='utf-8') as path_json:
        jsonx = json.load(path_json)
        with open(path_txt, 'w+') as ftxt:
            shapes = jsonx['shapes']
            # 获取图片长和宽
            width = jsonx['imageWidth']
            height = jsonx['imageHeight']

            for shape in shapes:
                # 获取矩形框两个角点坐标
                x1 = shape['points'][0][0]
                y1 = shape['points'][0][1]
                x2 = shape['points'][1][0]
                y2 = shape['points'][1][1]

                dw = 1. / width
                dh = 1. / height
                x = dw * (x1 + x2) / 2
                y = dh * (y1 + y2) / 2
                w = dw * abs(x2 - x1)
                h = dh * abs(y2 - y1)
                cat = shape['label']
                cat = classes2id[cat]
                yolo = f"{cat} {x} {y} {w} {h} \n"
                ftxt.writelines(yolo)


def handle_folder(dir_json, target_txt):
    """
    处理一个文件夹下所有的json文件
    :return:
    """
    os.makedirs(target_txt, exist_ok=True)
    list_json = os.listdir(dir_json)
    for cnt, json_name in enumerate(list_json):
        if os.path.splitext(json_name)[-1] == ".json":
            path_json = os.path.join(dir_json, json_name)
            path_txt = os.path.join(target_txt, json_name.replace('.json', '.txt'))
            json2txt(classes2id, path_json, path_txt)


if __name__ == "__main__":
    # ！！！！！注意：json标注的label要是数字，如果不是数字要修改代码
    # labelme标注的json标签文件目录和保存生成的txt标签的文件夹
    dir_json = r'F:\yolo_summary\yolo_code\car\521\train\tmp\tmp_json'
    target_txt = r'F:\yolo_summary\yolo_code\car\521\train\tmp\tmp_json_txt'
    classes2id = get_classes2id(dir_json)
    handle_folder(dir_json, target_txt)
